Flatten 2D vector

Time: O(1); Space: O(1); medium

Implement an iterator to flatten a 2D vector.

Example 1:

Input: vec2d =

[
  [1,2],
  [3],
  [4,5,6]
]

By calling next repeatedly until hasNext returns False, the order of elements returned by next should be Output: [1,2,3,4,5,6]

Hints:

  1. How many variables do you need to keep track? Two variables is all you need. Try with x and y.

  2. Beware of empty rows. It could be the first few rows.

  3. To write correct code, think about the invariant to maintain. What is it? The invariant is x and y must always point to a valid point in the 2d vector. Should you maintain your invariant ahead of time or right when you need it? Not sure? Think about how you would implement hasNext(). Which is more complex?

  4. Common logic in two different places should be refactored into a common method.

Follow up:

  • As an added challenge, try to code it using only iterators in Python.

[8]:

class Vector2D:
    x, y = 0, 0
    vec = None

    def __init__(self, vec2d):
        """
        :type vec2d: int[][]
        """
        self.vec = vec2d
        self.x = 0
        if self.x != len(self.vec):
            self.y = 0
            self.adjustNextIter()

    def next(self):
        """
        :rtype: int
        """
        ret = self.vec[self.x][self.y]
        self.y += 1
        self.adjustNextIter()
        return ret

    def hasNext(self):
        """
        :rtype: bool
        """
        return self.x != len(self.vec) and self.y != len(self.vec[self.x])

    def adjustNextIter(self):
        while self.x != len(self.vec) and self.y == len(self.vec[self.x]):
            self.x += 1
            if self.x != len(self.vec):
                self.y = 0

[9]:

L = [ [1,2], [3], [4,5,6] ]
v = Vector2D(L)
vl = []
while v.hasNext():
    vl.append(v.next())
assert vl == [1, 2, 3, 4, 5, 6]